import os
from tqdm import tqdm
from ..utils.processing_charls import CharlsProcessor
from ..build_embeddings import get_detailed_instruct
from .writer_encodings_embeddings import EncodingEmbeddingWriter


def process_encoding_embeddings_charls(
    codebook_path: str,
    output_dir: str,
    model_name: str,
    tensor_parallel_size: int,
    batch_size: int = 1024,
):
    """
    Processes embeddings for CHARLS encodings and saves them.
    """
    print("Loading CHARLS encodings...")
    data_processor = CharlsProcessor(codebook_path=codebook_path)
    encodings = data_processor.encodings
    print(f"Loaded {len(encodings)} CHARLS encodings.")

    encoding_items = list(encodings.items())
    texts_to_embed = [item[1].meaning for item in encoding_items]

    task = 'Generating representations of biomedical text'
    for i in tqdm(range(len(texts_to_embed)), desc="Adding detailed instruct"):
        texts_to_embed[i] = get_detailed_instruct(task, texts_to_embed[i])

    # Pair keys with processed texts
    encoding_items_with_processed_text = list(zip([item[0] for item in encoding_items], texts_to_embed))

    print("Writing embeddings for CHARLS encodings...")
    with EncodingEmbeddingWriter(
        output_dir,
        model_name,
        tensor_parallel_size=tensor_parallel_size
    ) as embedding_writer:
        for i in range(0, len(encoding_items_with_processed_text), batch_size):
            progress = (i + batch_size) / len(encoding_items_with_processed_text)
            print(f"Progress: {progress:.2%}")
            batch = encoding_items_with_processed_text[i:i + batch_size]
            embedding_writer.embedding_and_write_batch(batch)

    print("Finished processing all CHARLS encodings.")